import pathlib

import numpy

from problem1.checkpoint import Checkpoint
from utilities.checkpoint_manager import CheckpointManager


def main(checkpoint_manager: CheckpointManager,
         t_max: int = 500):
    t, checkpoint = checkpoint_manager.last_existed_loaded(Checkpoint.load)
    random = numpy.random.default_rng(1234)
    origin = numpy.zeros((10000, 2), dtype=numpy.int32)
    if checkpoint:
        random.__setstate__(checkpoint.random_state())
        caloric = checkpoint.calculation_state()
        print(f"t={t}: Loaded")
    else:
        caloric = numpy.copy(origin)
        Checkpoint(caloric, random.__getstate__()).save(checkpoint_manager.new(0))
        print(f"t={0}: Saved")

    for t in range(t + 1, t_max + 1):
        caloric += random.integers(-1, 1 + 1, size=caloric.shape)
        caloric = numpy.concatenate([caloric, origin])
        if t % 50 != 0:
            print(f"t={t}")
        else:
            Checkpoint(caloric, random.__getstate__()).save(checkpoint_manager.new(t))
            print(f"t={t}: Saved")


if __name__ == "__main__":
    main(CheckpointManager(pathlib.Path("./outputs")))
